#! /Pub/Users/wangyk/software/pixi/env/omicverse/.pixi/envs/ov/bin/python

def cpdb_v5_py310_v1(counts_file_path="your/path/anndata.h5ad",
                     cpdb_file_path='/Pub/Users/wangyk/project/Poroject/F241121002_20250210/out/1.sc/5.cellphone_db/v5.0.0/cellphonedb.zip',
                     meta_file_path="your/path/meta.data.tsv",
                     celltype_key='major_celltype',
                     ct1="Neutrophil",
                     ct2='.',
                     od='./',
                     interaction_gs: list|None =None,
                     figsize_of_dot: tuple = (7, 9),
                     figsize_of_heatmap: tuple = (6, 6),
                     figsize_of_chord: tuple = (5, 5),
                     **kwargs) -> None:
    """
    Run CellPhoneDB statistical analysis and generate visualizations.

    This function performs CellPhoneDB (CPDB) statistical analysis to infer cell-cell communication
    from single-cell RNA sequencing data and generates dot plot, heatmap, and chord diagram visualizations.

    Args:
        counts_file_path (str, optional): Path to the AnnData h5ad file containing count matrix.
            Defaults to "your/path/anndata.h5ad".
        cpdb_file_path (str, optional): Path to the CellphoneDB database zip file.
            Defaults to '/Pub/Users/wangyk/project/Poroject/F241121002_20250210/out/1.sc/5.cellphone_db/v5.0.0/cellphonedb.zip'.
        meta_file_path (str, optional): Path to the metadata tsv file defining barcodes to cell labels.
            Defaults to "your/path/meta.data.tsv".
            First column should be barcode, second column should be celltype.
        celltype_key (str, optional): Key in AnnData metadata or metadata file indicating cell types.
            Defaults to 'major_celltype'.
        ct1 (str, optional): Cell type 1 for dot plot and chord diagram visualization.
            Defaults to "Neutrophil".
        ct2 (str, optional): Cell type 2 for dot plot and chord diagram visualization.
            Defaults to '.', which means all other cell types will be considered.
        od (str, optional): Output directory to save results and plots.
            Defaults to './'.
        interaction_gs (list, optional): List of specific interactions (gene pairs) to highlight in the chord diagram.
            If None, top interactions related to ct1 (Neutrophil) will be automatically selected. Defaults to None.
        figsize_of_dot (tuple, optional): Figure size for the dot plot. Defaults to (7, 9).
        figsize_of_heatmap (tuple, optional): Figure size for the heatmap. Defaults to (6, 6).
        figsize_of_chord (tuple, optional): Figure size for the chord diagram. Defaults to (5, 5).
    """

    # 运行分析
    import pandas as pd
    from cellphonedb.src.core.methods import cpdb_statistical_analysis_method
    import os
    import pickle
    
    cpdb_results = cpdb_statistical_analysis_method.call(
        cpdb_file_path=cpdb_file_path,
        # mandatory: CellphoneDB database zip file.
        meta_file_path=meta_file_path,
        # mandatory: tsv file defining barcodes to cell label.
        counts_file_path=counts_file_path,
        # mandatory: normalized count matrix - a path to the counts file, or an in-memory AnnData   object
        counts_data='hgnc_symbol',
        # defines the gene annotation in counts matrix.
        # active_tfs_file_path = active_tf_path,
        # # optional: defines cell types and their active TFs.
        # microenvs_file_path = microenvs_file_path,
        # optional (default: None): defines cells per microenvironment.
        score_interactions=True,
        # optional: whether to score interactions or not.
        iterations=1,
        # denotes the number of shufflings performed in the analysis.
        threshold=0.05,
        # defines the min % of cells expressing a gene for this to be employed in the analysis.
        threads=5,
        # number of threads to use in the analysis.
        debug_seed=12,
        # debug randome seed. To disable >=0.
        result_precision=3,
        # Sets the rounding for the mean values in significan_means.
        pvalue=0.05,
        # P-value threshold to employ for significance.
        subsampling=False,
        # To enable subsampling the data (geometri sketching).
        subsampling_log=False,
        # (mandatory) enable subsampling log1p for non log-transformed data inputs.
        subsampling_num_pc=100,
        # Number of componets to subsample via geometric skectching (dafault: 100).
        subsampling_num_cells=1000,
        # Number of cells to subsample (integer) (default: 1/3 of the dataset).
        separator='|',
        # Sets the string to employ to separate cells in the results dataframes "cellA|CellB".
        debug=False,
        # Saves all intermediate tables employed during the analysis in pkl format.
        output_path=od,
        # Path to save results.
        output_suffix=None
        # Replaces the timestamp in the output files by a user defined string in the  (default:     None).
    )
    
    if not os.path.exists(od):
        os.makedirs(od, exist_ok=True, parents=True)

    with open(f'{od}/cpdb_results.pkl', 'wb') as file:
        pickle.dump(cpdb_results, file)

    import ktplotspy as kpy
    import matplotlib.pyplot as pl
    import scanpy as sc

    adata_ = sc.read_h5ad(counts_file_path)

    p = kpy.plot_cpdb(
        adata=adata_,
        cell_type1=ct1,
        cell_type2=ct2,
        means=cpdb_results['means'],
        pvals=cpdb_results['pvalues'],
        interaction_scores=cpdb_results['interaction_scores'],
        celltype_key=celltype_key,
        # genes = ["TGFB2", "CSF1R"],
        figsize=figsize_of_dot,
        title="Interactions between\nNeutrophil and other cells",
        max_size=4,
        highlight_size=0.75,
        degs_analysis=False,
        standard_scale=True,
        scale_alpha_by_interaction_scores=True,
    )
    p.save(filename = f"{od}/1.dotplot_of_Interactions.pdf") 

    pl.clf()
    kpy.plot_cpdb_heatmap(
        pvals=cpdb_results['pvalues'],
        figsize=figsize_of_heatmap,
        title="Sum of significant interactions",
        symmetrical=False,
    )
    pl.savefig(f"{od}/2.heatrmap_of_interactions.pdf",
               bbox_inches='tight', dpi=300)

    if interaction_gs is None:
        df = cpdb_results['significant_means']
        pipe_cols = [col for col in df.columns if '|' in col]
        no_pipe_cols = [col for col in df.columns if '|' not in col]
        id_vars = no_pipe_cols  # Columns to keep as identifiers
        df_long = pd.melt(df,
                          id_vars=id_vars,
                          value_vars=pipe_cols,
                          var_name='cell',
                          value_name='value')
        df_long = df_long[df_long['cell'].str.contains(
            'Neutrophil\|')].dropna().sort_values(by='value',    ascending=False)
        # gs_out = df_long[df_long['receptor_a'] == True].loc[:,'gene_a'].unique().tolist()
        gs_out = df_long.loc[:, 'gene_a'].unique().tolist()[0:5]
    else:
        gs_out = interaction_gs

    kpy.plot_cpdb_chord(
        adata=adata_,
        cell_type1=ct1,
        cell_type2=ct2,
        means=cpdb_results['means'],
        pvals=cpdb_results['pvalues'],
        deconvoluted=cpdb_results['deconvoluted'],
        celltype_key=celltype_key,
        interaction=gs_out,
        link_kwargs={"direction": 1, "allow_twist": True, "r1": 95, "r2": 90},
        sector_text_kwargs={"color": "black", "size": 14,
                            "r": 105, "adjust_rotation": True},
        legend_kwargs={"fontsize": 9,"loc":"lower center","bbox_to_anchor":(.98, -.02)},
        link_offset=1,
        figsize=figsize_of_chord,
    )
    pl.savefig(f"{od}/3.plot_cpdb_chord.pdf", bbox_inches='tight', dpi=300)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(
        description='Run CellPhoneDB analysis and generate visualizations.')
    parser.add_argument('--counts_file_path', type=str, default="your/path/anndata.h5ad",
                        help='Path to the AnnData h5ad file containing the count matrix, or path to the counts csv file, where the structure is rows = Genes x columns = cell_barcodes.', metavar='')
    parser.add_argument('--cpdb_file_path', type=str, default='/Pub/Users/wangyk/project/Poroject/F241121002_20250210/out/1.sc/5.cellphone_db/v5.0.0/cellphonedb.zip',
                        help='Path to the CellphoneDB database zip file.', metavar='')
    parser.add_argument('--meta_file_path', type=str, default="your/path/meta.data.tsv",
                        help='Path to the metadata tsv file defining barcodes to cell labels.', metavar='')
    parser.add_argument('--celltype_key', type=str, default='major_celltype',
                        help='Key in AnnData metadata or metadata file indicating cell types.', metavar='')
    parser.add_argument('--ct1', type=str, default="Neutrophil",
                        help='Cell type 1 for dot plot and chord diagram visualization.', metavar='')
    parser.add_argument('--ct2', type=str, default='.',
                        help='Cell type 2 for dot plot and chord diagram visualization. "." means all cell type', metavar='')
    parser.add_argument('--od', type=str, default='./',
                        help='Output directory to save results and plots.', metavar='')
    parser.add_argument('--interaction_gs', type=str, default=None,
                        help='List of specific interactions (gene pairs) to highlight in the chord diagram.', metavar='')
    parser.add_argument('--figsize_of_dot', type=int, nargs=2, default=(7, 9),
                        help='Figure size for the dot plot. eg: --figsize_of_dot 7 9', metavar='')
    parser.add_argument('--figsize_of_heatmap', type=int, nargs=2, default=(5, 5),
                        help='Figure size for the heatmap. eg: --figsize_of_heatmap 5 5', metavar='')
    parser.add_argument('--figsize_of_chord', type=int, nargs=2, default=(5, 5),
                        help='Figure size for the chord diagram. eg: --figsize_of_chord 5 5', metavar='')
    args = parser.parse_args()

    for i in ['figsize_of_dot', 'figsize_of_heatmap', 'figsize_of_chord']:
        args.__dict__[i] = tuple(args.__dict__[i])

    cpdb_v5_py310_v1(**vars(args))
